# 导入函数库
from jqdata import *


# 初始化函数，设定基准等等
def initialize(context):
    # 设定沪深300作为基准
    set_benchmark('000300.XSHG')
    # 开启动态复权模式(真实价格)
    set_option('use_real_price', True)
    # 输出内容到日志 log.info()
    log.info('初始函数开始运行且全局只运行一次')
    # 过滤掉order系列API产生的比error级别低的log
    log.set_level('order', 'error')

    ### 股票相关设定 ###
    # 股票类每笔交易时的手续费是：买入时佣金万分之三，卖出时佣金万分之三加千分之一印花税, 每笔交易佣金最低扣5块钱
    set_order_cost(OrderCost(close_tax=0.001, open_commission=0.00015, close_commission=0.00015, min_commission=5),
                   type='stock')

    ## 运行函数（reference_security为运行时间的参考标的；传入的标的只做种类区分，因此传入'000300.XSHG'或'510300.XSHG'是一样的）

    # 开盘时运行
    run_daily(trade, time='14:50')


def trade(context):
    # 先卖出股票腾仓位
    # 看看持有股票，如果有跌破20日线的，卖出
    long_positions_dict = context.portfolio.long_positions
    for position in list(long_positions_dict.values()):
        # 获取20日均价
        ma5 = get_bars(position.security, count=5, unit='1d', fields=['close'])['close'].mean()
        ma20 = get_bars(position.security, count=20, unit='1d', fields=['close'])['close'].mean()
        # 如果最新价格低于10日线，卖
        if position.price < ma5 or ma5 < ma20:
            order_target(position.security, 0)

    # 如果没有仓位了，今天不动了
    available_cash = context.portfolio.available_cash
    if available_cash < 3000:
        log.info("没钱了，今天不动了")
        return

    # 获取当日成交额前200只股票
    stocks = top_money_stock(context)
    # 账户总资产（现金+股票市值）, 每次买入一只票最多5成仓
    buy_max_per_stock = round(context.portfolio.total_value * 0.5, 2)
    # 循环处理每只备选票
    for stock in stocks:
        # 取得当前的现金
        cash = context.portfolio.available_cash
        if cash < 3000:
            log.info("没钱了，今天不继续买了")
            return
        # 获取20日均价
        ma5 = get_bars(stock, count=5, unit='1d', fields=['close'])['close'].mean()
        ma20 = get_bars(stock, count=20, unit='1d', fields=['close'])['close'].mean()
        # 最新价格
        current_price = get_bars(stock, count=1, unit='1m')['close'][0]
        # 最新价在20日线之上，且不持有，买入
        if (current_price > ma5) and (ma5 > ma20) and (context.portfolio.positions[stock].closeable_amount <= 0):
            # 用所有 cash 买入股票
            # order_value(stock, cash)
            # 用最多1/3仓位买入
            buy_cash = min(cash, buy_max_per_stock)
            order_value(stock, buy_cash)


def top_money_stock(context):
    # 返回成交额前200的股票
    # 获取所有上市股票
    all_stock = list(get_all_securities(['stock'], date=context.current_dt).index)
    # 获取所有股票当日成交额
    money = get_price(all_stock, context.current_dt, context.current_dt, fields=["money"], panel=False)
    # 排序取前200条
    sorted_money = money.sort_values(by='money', ascending=False).head(200)

    # 返回股票代码列表
    return list(sorted_money['code'])

